import os
from PIL import Image

# 输入和输出目录
input_dir = 'E:\postgraduate\大论文\img_out\YOLOv5识别'
txt_dir = 'E:\postgraduate\大论文数据集\分割数据集\IQI'
output_dir = 'E:\postgraduate\大论文\img_out\YOLOv5识别_裁剪'

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

# 获取所有图片文件
image_files = [f for f in os.listdir(input_dir) if f.endswith('.jpg') or f.endswith('.png')]
x = 1

for image_file in image_files:
    # 构建完整的图片和txt文件路径
    image_path = os.path.join(input_dir, image_file)
    txt_file_path = os.path.join(txt_dir, f'{os.path.splitext(image_file)[0]}.txt')

    # 检查对应的txt文件是否存在
    if not os.path.exists(txt_file_path):
        print(f"Warning: {txt_file_path} not found for {image_file}. Skipping.")
        continue

    # 读取xx1和xx2的值
    with open(txt_file_path, 'r') as f:
        xx1, xx2 = map(float, f.readline().strip().split())

    # 打开图片
    im = Image.open(image_path)

    # 计算裁剪后的区域
    width, height = im.size
    left = int(xx1)
    upper = 0
    right = int(xx2)
    lower = height

    # 裁剪图片
    cropped_im = im.crop((left, upper, right, lower))

    # 构造输出文件路径
    output_path = os.path.join(output_dir, image_file)

    # 保存裁剪后的图片
    cropped_im.save(output_path)
    x = x + 1
    print(x)

print("裁剪完成！")
